# 模型相关路由
import io
import os
import shutil
from pathlib import Path

from exts import db
from flask import request, jsonify, g, send_file
from flask_sse import sse
from models import User, Model, Plant, Dataset, Predict
import threading
from ACRmodel.train import train_model


def register_model_routes(app):
    # 查询所有模型
    @app.route('/plantACR/getModels', methods=['GET'])
    def get_models():
        resultVO = {
            "code": 200,
            "message": "成功获取模型列表",
            "data": {
                "models": []
            }
        }
        try:
            # 按ID升序查询所有模型
            models = Model.query.order_by(Model.id.asc()).all()
            if models:
                for model in models:
                    model_info = {
                        "id": model.id,
                        "model_name": model.model_name,
                        "user_id": model.user_id,
                        "plant_id": model.plant_id,
                        "dataset_id": model.dataset_id,
                        "user_name": model.user_name,
                        "plant_name": model.plant_name,
                        "dataset_name": model.dataset_name,
                        "description": model.description,
                        "epochs": model.epochs,
                        "stopped_epochs": model.stopped_epochs,
                        "public": model.public,
                        "status": model.status,
                        "accuracy": model.accuracy,
                        "auc": model.auc,
                        "precision": model.precision,
                        "recall": model.recall,
                        "f1": model.f1,
                        "metrics_image": model.metrics_image,
                        "confusion_matrix_image": model.confusion_matrix_image,
                        "file_size": model.file_size,
                        "create_time": model.create_time,
                        "update_time": model.update_time
                    }
                    resultVO["data"]["models"].append(model_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "模型列表为空"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取模型列表失败!"
            return jsonify(resultVO)

    # 查询公开模型
    @app.route('/plantACR/getPublicModels', methods=['GET'])
    def get_public_models():
        resultVO = {
            "code": 200,
            "message": "成功获取公开模型列表",
            "data": {
                "models": [],
                "a_models": []
            }
        }
        try:
            # 查询所有公开模型
            public_models = Model.query.filter_by(status='公开').order_by(Model.id.asc()).all()
            # 查询所有公开模型和自己的私密模型
            all_models = Model.query.filter(
                (Model.status == '公开') |
                ((Model.status == '私密') & (Model.user_id == g.id))
            ).order_by(Model.id.asc()).all()

            if all_models:
                # 将公开模型存入 models
                for model in public_models:
                    model_info = {
                        "id": model.id,
                        "model_name": model.model_name,
                        "user_id": model.user_id,
                        "plant_id": model.plant_id,
                        "dataset_id": model.dataset_id,
                        "user_name": model.user_name,
                        "plant_name": model.plant_name,
                        "dataset_name": model.dataset_name,
                        "description": model.description,
                        "epochs": model.epochs,
                        "stopped_epochs": model.stopped_epochs,
                        "public": model.public,
                        "status": model.status,
                        "accuracy": model.accuracy,
                        "auc": model.auc,
                        "precision": model.precision,
                        "recall": model.recall,
                        "f1": model.f1,
                        "metrics_image": model.metrics_image,
                        "confusion_matrix_image": model.confusion_matrix_image,
                        "file_size": model.file_size,
                        "create_time": model.create_time,
                        "update_time": model.update_time
                    }
                    resultVO["data"]["models"].append(model_info)

                # 将公开模型和自己的私密模型存入 a_models
                for model in all_models:
                    model_info = {
                        "id": model.id,
                        "model_name": model.model_name,
                        "user_id": model.user_id,
                        "plant_id": model.plant_id,
                        "dataset_id": model.dataset_id,
                        "user_name": model.user_name,
                        "plant_name": model.plant_name,
                        "dataset_name": model.dataset_name,
                        "description": model.description,
                        "epochs": model.epochs,
                        "stopped_epochs": model.stopped_epochs,
                        "public": model.public,
                        "status": model.status,
                        "accuracy": model.accuracy,
                        "auc": model.auc,
                        "precision": model.precision,
                        "recall": model.recall,
                        "f1": model.f1,
                        "metrics_image": model.metrics_image,
                        "confusion_matrix_image": model.confusion_matrix_image,
                        "file_size": model.file_size,
                        "create_time": model.create_time,
                        "update_time": model.update_time
                    }
                    resultVO["data"]["a_models"].append(model_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "公开模型列表为空"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取公开模型列表失败!"
            return jsonify(resultVO)

    # 查询我的模型
    @app.route('/plantACR/getMyModels', methods=['GET'])
    def get_my_models():
        resultVO = {
            "code": 200,
            "message": "成功获取个人模型列表",
            "data": {
                "models": []
            }
        }
        try:
            # 按ID升序查询所有模型
            models = Model.query.filter_by(user_id=g.id).order_by(Model.id.asc()).all()
            if models:
                for model in models:
                    model_info = {
                        "id": model.id,
                        "model_name": model.model_name,
                        "user_id": model.user_id,
                        "plant_id": model.plant_id,
                        "dataset_id": model.dataset_id,
                        "user_name": model.user_name,
                        "plant_name": model.plant_name,
                        "dataset_name": model.dataset_name,
                        "description": model.description,
                        "epochs": model.epochs,
                        "stopped_epochs": model.stopped_epochs,
                        "public": model.public,
                        "status": model.status,
                        "accuracy": model.accuracy,
                        "auc": model.auc,
                        "precision": model.precision,
                        "recall": model.recall,
                        "f1": model.f1,
                        "metrics_image": model.metrics_image,
                        "confusion_matrix_image": model.confusion_matrix_image,
                        "file_size": model.file_size,
                        "create_time": model.create_time,
                        "update_time": model.update_time
                    }
                    resultVO["data"]["models"].append(model_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "个人模型列表为空"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取个人模型列表失败!"
            return jsonify(resultVO)

    # 更新模型信息
    @app.route('/plantACR/updateModel', methods=['POST'])
    def update_model():
        resultVO = {
            "code": 200,
            "message": "模型信息更新成功",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            model_id = data.get('id')
            new_model_name = data.get('model_name')

            # 检查新模型名字是否已存在（自己除外）
            existing_model = Model.query.filter_by(model_name=new_model_name).first()
            if existing_model and existing_model.id != model_id:
                resultVO["code"] = 400
                resultVO["message"] = "设置的模型名称在数据库中已存在"
                return jsonify(resultVO)

            model = Model.query.get(model_id)
            if model:
                # 更新模型信息
                model.model_name = data.get('model_name', model.model_name)
                model.description = data.get('description', model.description)
                model.epochs = data.get('epochs', model.epochs)
                db.session.commit()
                resultVO["code"] = 200
                resultVO["message"] = "模型信息更新成功"
                resultVO["data"] = {
                    "model": {
                        "id": model.id,
                        "model_name": model.model_name,
                        "description": model.description,
                        "epochs": model.epochs
                    }
                }
                return jsonify(resultVO)
            else:
                resultVO["code"] = 404
                resultVO["message"] = "模型不存在"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，更新模型信息失败!"
            return jsonify(resultVO)

    # 更新模型公开状态
    @app.route('/plantACR/updateModelPublic', methods=['POST'])
    def update_model_public():
        resultVO = {
            "code": 200,
            "message": "模型公开状态更新成功",
            "data": {}
        }
        try:
            # 模型的状态：未训练, 训练中, 训练失败, 公开, 私密
            data = request.get_json()  # 从前端获取 JSON 数据
            model_id = data.get('id')
            model = Model.query.get(model_id)
            if model:
                if model.status in ["未训练", "训练中", "训练失败"]:
                    resultVO["code"] = 400
                    resultVO["message"] = "模型未训练，不能设置公开状态"
                    return jsonify(resultVO)

                # 公开状态取反
                model.public = not model.public
                model.status = '公开' if model.public else '私密'  # 状态同步
                db.session.commit()
                print("model.public:", model.public)
                print("model.status:", model.status)

                resultVO["code"] = 200
                resultVO["message"] = "模型公开状态更新成功"
                resultVO["data"] = {
                    "model": {
                        "id": model.id,
                        "public": model.public
                    }
                }
                return jsonify(resultVO)
            else:
                resultVO["code"] = 404
                resultVO["message"] = "模型不存在"
                return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，更新模型公开状态失败!"
            return jsonify(resultVO)

    # 删除模型
    @app.route('/plantACR/deleteModel', methods=['POST'])
    def delete_model():
        resultVO = {
            "code": 200,
            "message": "模型删除成功",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            model_id = data.get('id')
            model = Model.query.get(model_id)

            # 模型不存在
            if not model:
                resultVO["code"] = 404
                resultVO["message"] = "模型不存在"
                return jsonify(resultVO)

            # 检查权限，只有管理员或用户自己可以删除模型
            if g.role != 'admin' and g.id != model.user_id:
                resultVO["code"] = 403
                resultVO["message"] = "权限不足，无法删除该模型"
                return jsonify(resultVO)

            # 检查模型是否存在预测记录
            model_predicts = Predict.query.filter_by(model_id=model_id).first()
            if model_predicts:
                resultVO["code"] = 400
                resultVO["message"] = '该模型存在预测记录，无法删除'
                return jsonify(resultVO)

            # 删除本地文件 + 模型记录
            file_path = f"ACRmodel/static/weights/{model.plant_name}/{model.id}"
            if os.path.exists(file_path):
                shutil.rmtree(file_path)
            db.session.delete(model)
            db.session.commit()  # 提交删除操作

            resultVO["code"] = 200
            resultVO["message"] = "模型删除成功"
            resultVO["data"] = {
                "model": {
                    "id": model_id,
                    "model_name": model.model_name
                }
            }
            return jsonify(resultVO)
        except Exception as e:
            print("服务器错误：", e)
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，删除模型失败!"
            return jsonify(resultVO)

    # 创建数据集任务(未训练)
    @app.route('/plantACR/createModel', methods=['POST'])
    def create_model():
        resultVO = {
            "code": 200,
            "message": "已创建训练任务",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            user_id = g.id

            user = User.query.get(user_id)
            plant = Plant.query.filter_by(name=data.get('plant_name')).first()
            dataset = Dataset.query.filter_by(dataset_name=data.get('dataset_name')).first()
            existing_model = Model.query.filter_by(model_name=data.get('model_name')).first()
            # 检查用户 是否存在
            if not user:
                resultVO["code"] = 404
                resultVO["message"] = "登录信息异常，用户不存在"
                return jsonify(resultVO)

            # 检查植物 是否存在
            if not plant:
                resultVO["code"] = 404
                resultVO["message"] = "对应的植物在数据库中不存在：" + data.get('plant_name')
                return jsonify(resultVO)

            # 检查数据集 是否存在
            if not dataset:
                resultVO["code"] = 404
                resultVO["message"] = "对应的数据集在数据库中不存在：" + data.get('dataset_name')
                return jsonify(resultVO)

            # 检查模型名字 是否已存在
            if existing_model:
                resultVO["code"] = 400
                resultVO["message"] = "设置的模型名称在数据库中已存在，请更换名称"
                return jsonify(resultVO)

            # 创建新的模型记录
            new_model = Model(
                user_id=g.id,
                user_name=user.name,
                plant_id=plant.id,
                plant_name=plant.name,
                dataset_id=dataset.id,
                dataset_name=dataset.dataset_name,
                model_name=data.get('model_name'),
                description=data.get('description'),
                epochs=data.get('epochs', 50),  # 默认为50轮
                public=data.get('public', False),  # 公开, 私密, 默认私密（False）
                status="未训练",  # 模型的状态：未训练, 训练中, 训练失败, 公开, 私密
            )
            db.session.add(new_model)
            db.session.commit()
            resultVO["data"] = {
                "model": {
                    "id": new_model.id,
                    "model_name": new_model.model_name,
                    "user_name": new_model.user_name,
                    "plant_name": new_model.plant_name,
                    "dataset_name": new_model.dataset_name,
                    "public": new_model.public,
                    "status": new_model.status,
                    "description": new_model.description,
                    "epochs": new_model.epochs,
                    "create_time": new_model.create_time
                }
            }
            return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            print("模型训练任务添加失败:", e)
            resultVO["code"] = 500
            resultVO["message"] = "模型训练任务添加失败!"
            return jsonify(resultVO)

    # 开始训练模型
    @app.route('/plantACR/startTrain', methods=['POST'])
    def start_train_model():
        resultVO = {
            "code": 200,
            "message": "模型训练开始",
            "data": {}
        }
        data = request.get_json()  # 从前端获取 JSON 数据
        model_id = data.get('id')
        try:
            # 检查模型是否存在
            model = Model.query.get(model_id)
            if not model:
                resultVO["code"] = 404
                resultVO["message"] = "模型不存在"
                return jsonify(resultVO)

            # 检查权限，只有管理员或用户自己能训练模型
            if g.id != model.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = "权限不足，无法训练该模型"
                return jsonify(resultVO)

            # 检查是否有其他模型训练任务正在进行
            training_model = Model.query.filter_by(status="训练中").first()
            if training_model:
                resultVO["code"] = 400
                resultVO["message"] = f"模型 {training_model.model_name} 正在训练中，请等待完成"
                return jsonify(resultVO)

            # 检查模型状态是否可以训练，模型的状态：未训练, 训练中, 训练失败, 公开, 私密
            if model.status in ('未训练', '训练失败'):
                # 启动训练任务
                model.status = "训练中"
                epochs = data.get('epochs', model.epochs)  # 默认为数据库中的值
                model.epochs = epochs
                db.session.commit()
                resultVO["code"] = 200
                resultVO["message"] = "模型开始训练"
                resultVO["data"] = {
                    "model": {
                        "id": model.id,
                        "model_name": model.model_name,
                        "status": model.status,
                        "epochs": model.epochs
                    }
                }
                # 启动后台线程执行训练
                threading.Thread(target=start_train, args=(model_id, epochs)).start()
                return jsonify(resultVO)
            # 训练中、公开、私密状态是无法发送训练请求的，返回数据异常
            else:
                resultVO["code"] = 400
                resultVO["message"] = "数据异常！"
                return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            model = Model.query.get(model_id)
            if model:
                model.status = "训练失败"
                db.session.commit()
            print("模型训练任务启动失败:", str(e))
            print(e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，模型训练失败!"
            return jsonify(resultVO)

    # 在后台线程中执行训练
    def start_train(model_id, epochs):
        """
            status: 未训练, 训练中, 训练失败, 公开, 私密
                训练完成后，更新模型的状态

            :param
                model_id: 模型id
                epochs: 训练轮数
            :return
                1. 训练成功：更新模型状态为公开/私密，发送SSE消息
                2. 训练失败：更新模型状态为训练失败，发送SSE消息
                3. 服务器错误：更新模型状态为训练失败，发送SSE消息
                4. 服务器错误：更新模型状态为训练失败，发送SSE消息
        """
        with app.app_context():
            train_result = {}
            try:
                model = Model.query.get(model_id)
                if model:
                    print("启动模型训练...")
                    sse_model_status(model_id, "训练中", train_result, "模型训练开始")
                    # 调用训练函数启动训练
                    train_result = train_model(plant_name=model.plant_name, dataset_id=model.dataset_id,
                                               model_id=model.id,
                                               epochs=epochs)

                    if train_result["status"] == "训练成功":
                        print("训练成功！模型评估:", train_result)
                        status = '公开' if model.public else '私密'
                        model = Model.query.get(model_id)
                        model.status = status
                        model.accuracy = train_result["accuracy"]
                        model.auc = train_result["auc"]
                        model.precision = train_result["precision"]
                        model.recall = train_result["recall"]
                        model.f1 = train_result["f1"]
                        model.metrics_image = train_result["metrics_image"]
                        model.confusion_matrix_image = train_result["confusion_matrix_image"]
                        model.file_size = train_result["file_size"]
                        model.stopped_epochs = train_result["stopped_epochs"]
                        db.session.commit()
                        sse_model_status(model_id, status, train_result, train_result["message"])
                    else:
                        print("训练失败：", train_result)
                        sse_model_status(model_id, "训练失败", train_result, train_result["message"])
                else:
                    print("模型训练任务不存在")
                    sse_model_status(model_id, "训练失败", train_result, "模型训练任务不存在")
            except Exception as e:
                db.session.rollback()
                print("模型训练失败:", e)
                sse_model_status(model_id, "训练失败", train_result, "服务器错误，训练失败！")
            finally:
                print("finally:当前模型训练结束")

    # 发送SSE，并更新模型在数据库中的状态，模型的状态：未训练, 训练中, 训练失败, 公开, 私密
    def sse_model_status(model_id, status, result, sse_message=None):
        model = Model.query.get(model_id)
        if model:
            model.status = status
            db.session.commit()
            # 这里向所有订阅了 train_status 事件的客户端（前端）发送一条消息
            try:
                # # 发送SSE消息
                sse.publish({
                    "model_id": model_id,
                    "model_name": model.model_name,
                    "status": status,
                    "message": sse_message,
                    "result": result
                }, type='train_status', channel='sse')
                print("SSE消息已发布")
            except Exception as e:
                print("SSE发布失败：", e)
            return "消息已发送"
        else:
            print("模型不存在")
            sse.publish({
                "model_id": model_id,
                "model_name": "",
                "status": status,
                "message": sse_message,
                "result": result
            }, type='train_status', channel='sse')
            return "消息已发送，模型不存在"

    # 下载模型
    @app.route('/plantACR/downloadModel/<int:model_id>', methods=['GET'])
    def download_model(model_id):
        resultVO = {
            "code": 200,
            "message": ""
        }
        try:
            model = Model.query.get(model_id)
            if not model:
                resultVO["code"] = 404
                resultVO["message"] = "数据集不存在"
                return jsonify(resultVO)  # 虽然是json，但是前端以二进制文件Blob接受，需要进一步读取

            # 公开数据集都可以下载，私密数据集只有管理员和用户自己可以下载
            if (not model.public) and g.id != model.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = '权限不足！'
                return jsonify(resultVO)

            # 文件路径
            file_name = f"model_{model_id}.zip"
            dataset_file = Path(app.root_path) / "ACRmodel" / "static" / "weights" / model.plant_name / str(
                model_id) / file_name

            # 文件不存在
            if not dataset_file.exists():
                resultVO["code"] = 404
                resultVO["message"] = "文件不存在"
                return jsonify(resultVO)

            # 获取二进制流
            with open(dataset_file, 'rb') as file:
                bytes_io = io.BytesIO(file.read())
            # 返回文件给前端
            return send_file(bytes_io,
                             mimetype='application/zip, application/octet-stream',
                             as_attachment=True,  # 强制浏览器以附件形式下载文件，而非直接打开
                             download_name=file_name  # 设置默认下载文件名
                             )
        except Exception as e:
            print(str(e))
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，下载失败！"
            return jsonify(resultVO)

    return app
